﻿// RootRpcClient.cpp : Defines the exported functions for the DLL application.
//

#include "pch.h"

#include "..\RootRpcInterface\RootRpcInterface.h"
#include "RootRpcClient.h"

ROOTRPC_REFERENCE reference = nullptr;
WINERR initStatus = RPC_S_NO_BINDINGS;

ROOTRPC_EXPORTED WINERR InitRootRpc()
{
	RPC_STATUS status = NO_ERROR;
	// Start the service
	WINERR err = NO_ERROR;
	SC_HANDLE schandle = nullptr;
	SC_HANDLE servicehandle = nullptr;
	// First make sure we haven't already initialized
	if (reference)
	{
		return ERROR_ALREADY_INITIALIZED;
	}
	// Get the service's SC_HANDLE is available
	schandle = ::OpenSCManagerW(nullptr, nullptr, SC_MANAGER_CONNECT);
	if (!schandle)
	{
		err = ::GetLastError();
		goto Cleanup;
	}
	servicehandle = ::OpenServiceW(
		schandle,
		SERVICE_NAME,
		SERVICE_QUERY_STATUS | SERVICE_START);
	if (!servicehandle)
	{
		err = ::GetLastError();
		goto Cleanup;
	}
	// OK, we have the service HANDLE. Don't bother to check the status, just go.
	if (!::StartServiceW(servicehandle, 0x0, nullptr))
	{
		err = ::GetLastError();
		if (ERROR_SERVICE_ALREADY_RUNNING == err)
		{
			// If it's already running, that's great, we can get on with life.
			err = NO_ERROR;
		}
		else
		{
			goto Cleanup;
		}
	}
	// Server should be running now. OK, let's bind it.
	status = ::RpcBindingFromStringBindingW(
		(RPC_WSTR)L"ncalrpc:[RootRpc,Security=impersonation dynamic false]",
		&RootRpcHandle);
	if (status)
	{
		goto Cleanup;
	}
	// Server should be bound and ready, so grab a reference context (also serves as a test).
	RpcTryExcept
		for (int i = 0; i < 20; i++)
		{	// Might take a few attempts
			status = ::RpcMgmtIsServerListening(RootRpcHandle);
			if (status)
			{
				Sleep(50);
			}
			else
			{
				break;
			}
		}
		// Should now be properly connected.
		err = ::RootRpcInit(&reference);
	RpcExcept(EXCEPTION_EXECUTE_HANDLER)
		status = ::RpcExceptionCode();
	RpcEndExcept
Cleanup:
	if (schandle)
	{
		::CloseServiceHandle(schandle);
		schandle = nullptr;
	}
	if (servicehandle)
	{
		::CloseServiceHandle(servicehandle);
		servicehandle = nullptr;
	}

	if (err)
	{
		// The problem was probably in starting the service.
		return err;
	}
	// Either the RPC server is responding, or it's not but the service (was already?) started.
	return status;
}

ROOTRPC_EXPORTED WINERR CloseRootRpc()
{
	WINERR err = NO_ERROR;
	// This is extremely simple compared to the other function...
	if (!reference)
	{
		return RPC_S_NO_BINDINGS;
	}
	RpcTryExcept
		err = ::RootRpcRelease(&reference);
	RpcExcept(EXCEPTION_EXECUTE_HANDLER)
		err = ::RpcExceptionCode();
	RpcEndExcept
	::RpcBindingFree(&RootRpcHandle);
	return err;
}

ROOTRPC_EXPORTED WINERR TestRootRpc()
{
	RPC_STATUS status;
	// Verify that the server is available and that calls work
	RpcTryExcept
		status = ::RpcMgmtIsServerListening(RootRpcHandle);
		if (!status)
		{
			::RootRpcServerVersion(nullptr, nullptr, nullptr);
		}
	RpcExcept(EXCEPTION_EXECUTE_HANDLER)
		return ::RpcExceptionCode();
	RpcEndExcept

	return status;
}

ROOTRPC_EXPORTED WINERR GetRootRpcInitStatus()
{
	return initStatus;
}

ROOTRPC_EXPORTED DWORD RpcClientVersion(
	PDWORD major,
	PDWORD minor,
	PDWORD bugfix)
{
	if (major) *major = VER_MAJOR;
	if (minor) *minor = VER_MINOR;
	if (bugfix) *bugfix = VER_UPDATE;
	return VER_BUILD;
}

ROOTRPC_EXPORTED DWORD RootServerVersion(
	PDWORD major,
	PDWORD minor,
	PDWORD bugfix)
{
	RpcTryExcept
		return ::RootRpcServerVersion(major, minor, bugfix);
	RpcExcept(EXCEPTION_EXECUTE_HANDLER)
		::SetLastError(::RpcExceptionCode());
		return MAXDWORD;
	RpcEndExcept
}

ROOTRPC_EXPORTED HANDLE RootGetServerProcess(
	)
{
	HANDLE ret = nullptr;
	WINERR err;
	RpcTryExcept
		err = ::RootRpcGetServerProcess(&ret);
	RpcExcept(EXCEPTION_EXECUTE_HANDLER)
		err = ::RpcExceptionCode();
	RpcEndExcept
	if (err)
	{
		::SetLastError(err);
	}
	return ret;
}

ROOTRPC_EXPORTED HANDLE RootGetServerToken(
	)
{
	HANDLE ret = nullptr;
	WINERR err;
	RpcTryExcept
		err = ::RootRpcGetServerToken(&ret);
	RpcExcept(EXCEPTION_EXECUTE_HANDLER)
		err = ::RpcExceptionCode();
	RpcEndExcept
	if (err)
	{
		::SetLastError(err);
	}
	return ret;
}

void __RPC_FAR * __RPC_USER midl_user_allocate(size_t len)
{
	return(malloc(len));
}

void __RPC_USER midl_user_free(void __RPC_FAR * ptr)
{
	free(ptr);
}

void __RPC_USER HANDLE_to_xmit(HANDLE * handle, RPC_PTR * * transmitted)
{
	// Server will duplicate this HANDLE from us
	// However, we need to allocate memory for its transmission
	*transmitted = reinterpret_cast<RPC_PTR*>(midl_user_allocate(sizeof(RPC_PTR)));
	**transmitted = reinterpret_cast<RPC_PTR>(*handle);
}

void __RPC_USER HANDLE_from_xmit(RPC_PTR * transmitted, HANDLE * handle)
{
	// Server duplicated this HANDLE; value is already a valid HANDLE for us
	*handle = reinterpret_cast<HANDLE>(*transmitted);
}

void __RPC_USER HANDLE_free_inst( HANDLE * handle)
{
	//	No-op; client will manually close handles
	UNREFERENCED_PARAMETER(handle);
}

void __RPC_USER HANDLE_free_xmit( RPC_PTR  * transmitted)
{
	// Don't have to do anything but free the buffer for the transmitted value
	midl_user_free(transmitted);
}

void __RPC_USER CHANGEHANDLE_to_xmit(CHANGEHANDLE * handle, RPC_PTR * * transmitted)
{
	// Server will duplicate this HANDLE from us
	// However, we need to allocate memory for its transmission
	*transmitted = reinterpret_cast<RPC_PTR*>(midl_user_allocate(sizeof(RPC_PTR)));
	**transmitted = reinterpret_cast<RPC_PTR>(*handle);
}

void __RPC_USER CHANGEHANDLE_from_xmit(RPC_PTR * transmitted, CHANGEHANDLE * handle)
{
	// Server duplicated this HANDLE; value is already a valid HANDLE for us
	*handle = reinterpret_cast<CHANGEHANDLE>(*transmitted);
}

void __RPC_USER CHANGEHANDLE_free_inst( CHANGEHANDLE * handle)
{
	//	No-op; client will manually close handles
	UNREFERENCED_PARAMETER(handle);
}

void __RPC_USER CHANGEHANDLE_free_xmit( RPC_PTR  * transmitted)
{
	// Don't have to do anything but free the buffer for the transmitted value
	midl_user_free(transmitted);
}

void __RPC_USER HKEY_to_xmit(HKEY * handle, RPC_PTR * * transmitted)
{
	// Server will duplicate this HANDLE from us
	// However, we need to allocate memory for its transmission
	*transmitted = reinterpret_cast<RPC_PTR*>(midl_user_allocate(sizeof(RPC_PTR)));
	**transmitted = reinterpret_cast<RPC_PTR>(*handle);
}

void __RPC_USER HKEY_from_xmit(RPC_PTR * transmitted, HKEY * handle)
{
	// Server duplicated this HANDLE; value is already a valid HANDLE for us
	*handle = reinterpret_cast<HKEY>(*transmitted);
}

void __RPC_USER HKEY_free_inst(HKEY * handle)
{
	//	No-op; client will manually close handles
	UNREFERENCED_PARAMETER(handle);
}

void __RPC_USER HKEY_free_xmit(RPC_PTR  * transmitted)
{
	// Don't have to do anything but free the buffer for the transmitted value
	midl_user_free(transmitted);
}

LONG MapHKeyForRpc(PHKEY key)
{
	// Only initialize the HKCR and HKCU values when needed
	static HKEY hkcr = nullptr;
	static HKEY hkcu = nullptr;
	LONG status;
	switch (reinterpret_cast<DWORD>(*key))
	{
	case (DWORD)HKEY_CLASSES_ROOT:
		if (!hkcr)
		{
//			status = ::RegCreateKeyExW(HKEY_CLASSES_ROOT, L"", 0x0, nullptr, 0x0, KEY_READ, nullptr, &hkcr, nullptr);
			// Using the self-opened key was failing for some reason. Let's try with the server-supplied key...
			HANDLE token = nullptr;
			if (!::OpenProcessToken(::GetCurrentProcess(), TOKEN_ALL_ACCESS, &token))
			{
				// Can't get our own process token? Well, that's awkward...
				*key = nullptr;
				return ::GetLastError();
			}
			status = ::RootRpcRegOpenUserClassesRoot(token, 0x0, KEY_ALL_ACCESS, &hkcr);
			if (!hkcr)
			{
				*key = nullptr;
				::SetLastError(status);
				return status;
			}
		}
		// OK, HKCR should be set now.
		*key = hkcr;
		return NO_ERROR;
	case (DWORD)HKEY_CURRENT_USER:
		if (!hkcu)
		{
//			status = ::RegCreateKeyExW(HKEY_CURRENT_USER, L"", 0x0, nullptr, 0x0, KEY_READ, nullptr, &hkcr, nullptr);
			// Again, let's use the server-supplied version.
			status = ::RootRpcRegOpenCurrentUser(KEY_ALL_ACCESS, &hkcu);
			if (!hkcu)
			{
				*key = nullptr;
				return status;
			}
		}
		// OK, HKCU should be set now.
		*key = hkcu;
		return NO_ERROR;
	default:
		return NO_ERROR;
	}
}

DWORD GetMultiStringSize(PWCHAR msz)
{
	DWORD ret = 1;
	if (!msz)
	{
		return 0x0;
	}
	while (msz[ret] || msz[ret - 1])
	{
		ret++;
	}
	// ret now contains the index of the last null byte
	return (ret + 1);
}